import json
import time
from enum import Enum
from typing import Dict, List, Optional, Union

from pydantic import Field

from app.agent.base import BaseAgent
from app.flow.base import BaseFlow
from app.llm import LLM
from app.logger import logger
from app.schema import AgentState, Message, ToolChoice
from app.tool import PlanningTool

import re


class PlanStepStatus(str, Enum):
    """Enum class defining possible statuses of a plan step"""

    NOT_STARTED = "not_started"
    IN_PROGRESS = "in_progress"
    COMPLETED = "completed"
    BLOCKED = "blocked"

    @classmethod
    def get_all_statuses(cls) -> list[str]:
        """Return a list of all possible step status values"""
        return [status.value for status in cls]

    @classmethod
    def get_active_statuses(cls) -> list[str]:
        """Return a list of values representing active statuses (not started or in progress)"""
        return [cls.NOT_STARTED.value, cls.IN_PROGRESS.value]

    @classmethod
    def get_status_marks(cls) -> Dict[str, str]:
        """Return a mapping of statuses to their marker symbols"""
        return {
            cls.COMPLETED.value: "[✓]",
            cls.IN_PROGRESS.value: "[→]",
            cls.BLOCKED.value: "[!]",
            cls.NOT_STARTED.value: "[ ]",
        }


class FPlanningFlow(BaseFlow):
    """A flow that manages planning and execution of tasks using agents."""

    llm: LLM = Field(default_factory=lambda: LLM())
    planning_tool: PlanningTool = Field(default_factory=PlanningTool)
    executor_keys: List[str] = Field(default_factory=list)
    active_plan_id: str = Field(default_factory=lambda: f"plan_{int(time.time())}")
    current_step_index: Optional[int] = 0
    before_message: Optional[str] = None
    step_infomation: Optional[str] = None
    current_step_infomation: Optional[str] = None
    goal: Optional['str'] = None
    current_step_agent: Optional['str'] = None

    def __init__(
            self, agents: Union[BaseAgent, List[BaseAgent], Dict[str, BaseAgent]], **data
    ):
        # Set executor keys before super().__init__
        if "executors" in data:
            data["executor_keys"] = data.pop("executors")

        # Set plan ID if provided
        if "plan_id" in data:
            data["active_plan_id"] = data.pop("plan_id")

        # Initialize the planning tool if not provided
        if "planning_tool" not in data:
            planning_tool = PlanningTool()
            data["planning_tool"] = planning_tool

        # Call parent's init with the processed data
        super().__init__(agents, **data)

        # Set executor_keys to all agent keys if not specified
        if not self.executor_keys:
            self.executor_keys = list(self.agents.keys())

    def get_executor(self, step_type: Optional[str] = None) -> BaseAgent:
        """
        Get an appropriate executor agent for the current step.
        Can be extended to select agents based on step type/requirements.
        """
        import re
        for key in step_type.keys():
            stpe_key = step_type[key]
            match = re.match(r"^\s*([^:]+?)\s*:", stpe_key)
            if match:
                step_type = match.group(1)

        # If step type is provided and matches an agent key, use that agent
        if step_type and step_type in self.agents:
            return self.agents[step_type]

        # Otherwise use the first available executor or fall back to primary agent
        for key in self.executor_keys:
            if key in self.agents:
                return self.agents[key]

        # Fallback to primary agent
        return self.primary_agent

    async def execute(self, input_text: str) -> str:
        """Execute the planning flow with agents."""
        try:
            if not self.primary_agent:
                raise ValueError("No primary agent available")

            # Create initial plan if input provided
            if input_text:
                await self._create_initial_plan(input_text)

            result = ""
            while True:
                # Get current step to execute
                self.current_step_agent, step_info = await self._get_current_step_info()

                """
                # Exit if no more steps or plan completed
                if self.current_step_index is None:
                    result += await self._finalize_plan()
                    break
                """

                if self.current_step_index is None:
                    break

                executor = self.agents[self.current_step_agent] if self.agents[
                    self.current_step_agent] else self.primary_agent
                step_result = await self._execute_step(executor, step_info)
                result += step_result + "\n"
                self.before_message = result

            return result

        except Exception as e:
            logger.error(f"Error in PlanningFlow: {str(e)}")
            return f"Execution failed: {str(e)}"

    async def _create_initial_plan(self, request: str) -> None:
        """Create an initial plan based on the request using the flow's LLM and PlanningTool."""
        logger.info(f"Creating initial plan with ID: {self.active_plan_id}")

        # Create a system message for plan creation
        system_message = Message.system_message(
            "1. 分析需求以理解任务范围,创建任务计划,标题,以及单步计划等"
            "2. 单步计划的具体内容需要根据Agent的具体能力进行规划,单步计划若存在前后关联性,应考虑前一步计划的结果能否满足下一步的需求"
            "3. 使用 'Agnet名称:计划' 结构把拆分的单步计划输出"
            "4. 可供选择的Agent包括以下内容："
            "  - manus：这是擅长调用各种工具的Agent,包括打开浏览器,打开文档,读取数据等"
            "  - cot：这是擅长数据分析的Agent,但无法调用工具"
            "  - score：这是专用于简历匹配以及打分的Agent,但无法调用工具"
            "5. 输出的格式：Goal："";Steps:1.Agnet1:计划1;2.Agnet2:计划2]"
            "6. 仅需要输出上述格式的内容,且输出需要遵守格式规定"
        )

        # Create a user message with the request
        user_message = Message.user_message(
            f"创建一个分步计划解决需求: {request}"
        )

        response = await self.llm.ask(
            messages=[user_message],
            system_msgs=[system_message],
            stream=False
        )
        print(response)
        self.step_infomation = response

    async def _get_current_step_info(self) -> tuple[Optional[str], Optional[str]]:

        if self.step_infomation:
            goal_part, steps_part = self.step_infomation.split("Steps:", 1)

            matches = re.findall(r"(\d+)\.\s*([^:]+?)\s*:\s*(.*?)(?=\s*\d+\.|\s*$)", steps_part)


            if self.current_step_index == len(matches):

                self.current_step_index = None

                return None,None

            if self.current_step_index is not None:

                self.current_step_index += 1

                return matches[self.current_step_index - 1][1], matches[self.current_step_index - 1][2]



    async def _execute_step(self, executor: BaseAgent, step_info: dict) -> str:
        """Execute the current step with the specified agent using agent.run()."""

        step_prompt = f"""

        YOUR CURRENT TASK:
        You are now working on step {self.current_step_index}: "{step_info}"

        Please execute this step using the appropriate tools. When you're done, provide a summary of what you accomplished.
        """

        if self.before_message:
            step_prompt = 'before_message_begin ' + self.before_message + ' before_message_end' + step_prompt

        # Use agent.run() to execute the step
        try:
            step_result = await executor.run(step_prompt)
            return step_result

        except Exception as e:
            logger.error(f"Error executing step {self.current_step_index}: {e}")
            return f"Error executing step {self.current_step_index}: {str(e)}"

    async def _finalize_plan(self) -> str:
        """Finalize the plan and provide a summary using the flow's LLM directly."""

        agent = self.primary_agent
        summary_prompt = f"""
        The plan {self.goal} has been completed. 
        Please provide a summary of what was accomplished and any final thoughts.
        """
        summary = await agent.run(summary_prompt)

        return summary

